#include <iostream>
#include <iomanip>
#include <vector>
#include <array>
#include <string>
#include <sstream>
#include <numeric>
#include <cstring>
#include "Core/VecCompare.h"
#include "Core/TensorExpr.h"
#include "FiniteDiff/Multigrid.h"
#include "Embedded/EmbeddedGrid.h"
#include "Embedded/EmLevelOp.h"
#include "Embedded/EmFuncFiller.h"
#include "Embedded/EmIntergrid.h"
#include "Embedded/EMixedDer.h"
#include "Embedded/BlockRelaxation.h"
#include "../example/common/RotatedBox.h"
#include "../example/common/TestUtility.h"

#if DIM == 2

template class Tensor<Real,2>;

const char *bcTypes = "DDDD";
const Real theta = M_PI / 6;
Vec<Real, 3> B { 1.0, 0.0, 2.0 };
Vec<Real, 3> A;

class ModifiedMixedDerOp : public MGLevelOp<2>
{
public:
  enum { Dim = 2 };
  using EmGrid = EmbeddedGrid<Dim>;

  ModifiedMixedDerOp(const EmGrid &aGrid, const Vec<Real, 3> &aA, Real w4j, const char *aBcTypes)
      : MGLevelOp<2>(aGrid, w4j, aBcTypes),
        emGrid(aGrid),
        eff(aGrid),
        elv(aGrid),
        emd(aGrid, aA),
        br(aGrid, emd.getSparseOp(), 1e10)
  {
  }

  ~ModifiedMixedDerOp() = default;

  Real computeNorm(const Tensor<Real, Dim> &phi) const
  {
    return elv.NRM(phi, CellCentered, 0);
  }

  void fillGhosts(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, const Tensor<Real, Dim-1> *bcData) const
  {
    MGLevelOp<2>::fillGhosts(phi, rhs, bcData);
    eff.transferBdryValues(phi, bcData[2*Dim]);
    br.apply(phi, rhs);
  }

  void fillGhosts(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, int) const
  {
    MGLevelOp<2>::fillGhosts(phi, rhs, 0);
    eff.transferBdryValues(phi, 0);
    br.apply(phi, rhs);
  }

  void applySmoother(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &smoothed, const Tensor<Real, Dim> &rhs) const
  {
    emd.relax(phi, rhs, smoothed, weightForJacobi);
  }

  void computeResidual(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &rsd, const Tensor<Real, Dim> &rhs) const
  {
    emd.apply(phi, rhs, rsd);
  }

  const EmGrid &getEmGrid() const { return emGrid; }

  const EMixedDer &getMixedDerOp() const { return emd; }

  using MGLevelOp<2>::getDomain;

protected:
  const EmGrid &emGrid;
  EmFuncFiller<Dim> eff;
  EmLevelOp<Dim> elv;
  EMixedDer emd;
  BlockRelaxation<Dim> br;
  using MGLevelOp<2>::weightForJacobi;
};

const ModifiedMixedDerOp *pDbgOp;

template <class T1, class T2>
std::array<Real,3> testMixedDer(const std::vector<const MGLevelOp<2> *> &opHier,
                                const std::vector<const Intergrid<2> *> &trHier,
                                const T1 &PHI, const T2 &D2PHI)
{
  const int Dim = 2;
  const EmbeddedGrid<Dim> &emGrid = static_cast<const ModifiedMixedDerOp *>(opHier[0])->getEmGrid();
  EmLevelOp<Dim> elv(emGrid);
  EmFuncFiller<Dim> eff(emGrid);
  const Box<Dim> gbx = emGrid.getGhostedBox();
  Tensor<Real, Dim> phi(gbx);
  Tensor<Real, Dim> err(gbx);
  Tensor<Real, Dim> rhs(gbx);
  Tensor<Real, Dim-1> cData[2*Dim+1];
  eff.fillPointValues(err, CellCentered, PHI);
  eff.fillPointValues(rhs, CellCentered, D2PHI);
  phi = 0.0;
  for(int d = 0; d < Dim; ++d) {
    for (int m : {0, 1}) {
      cData[2 * d + m].resize(reduce(gbx, d));
      cData[2*d+m] = 0.0;
    }
  }
  cData[2*Dim].resize(emGrid.bdryNodes.size());
  eff.fillBdryValues(cData[2*Dim], "curvilinear", PHI);

  Multigrid<Dim> mg(opHier, trHier);
  mg.setParam({3, 3, 10}, 30, 0.0, 1e-13);
  CPUTimer timer;
  timer.reset();
  mg.solve(phi, rhs, cData);
  std::cout << "Multigrid : " << std::setprecision(5) << timer() << "s. " << std::endl << std::defaultfloat;
  elv.zeroExtNodes(phi);
  err = err - phi;
  elv.zeroExtNodes(err);
  std::array<Real, 3> errnorm;
  for(int p = 0; p <= 2; ++p)
    errnorm[p] = elv.NRM(err, CellCentered, p);
  return errnorm;
}

template <class T1, class T2>
void doTest(const std::vector<int> &gridSize, int numCompHier, const T1 &PHI, const T2 &D2PHI)
{
  const int Dim = 2;
  const Real tol = 1e-12;
  using rVec = Vec<Real, Dim>;
  rVec basis[2] {
      { cos(theta), sin(theta) },
      { -sin(theta), cos(theta) }
  };
  rVec translate {sin(theta), 0.0 };
  const int Nx = gridSize.back();
  translate = translate + 1.0 / Nx;
  const int Lx = std::ceil((cos(theta) + sin(theta)) * Nx) + 2;
  RotatedBox<Dim> rb(1.0, basis);
  const int numGrid = gridSize.size();
  std::vector<const EmbeddedGrid<Dim> *> gridHier;
  std::vector<const ModifiedMixedDerOp *> opHier;
  std::vector<const EmIntergrid<Dim> *> trHier;
  for(int n = 0; n < numGrid; ++n) {
    const int N = gridSize[n];
    std::cout << "\n\nInitializing h = " << oneOver(N) << std::endl;
    const int L = Lx << (numGrid - n - 1);
    RectDomain<Dim> rd(Box<Dim>(0, L-1), 1.0/N, CellCentered, 2);
    EmbeddedGrid<Dim> *pGrid = new EmbeddedGrid<Dim>(
        rd,
        [&](rVec q, rVec &closest, rVec &normal, bool &contained, Real tol) {
          q = q - translate;
          rb.check(q, closest, normal, contained, tol);
          closest = closest + translate;
        }, false, 0.15, tol);
//    if(get_dbglevel() >= 2)
//      pGrid->dump(tmpos);
    gridHier.push_back(pGrid);
    ModifiedMixedDerOp *pOp = new ModifiedMixedDerOp(*pGrid, A, 0.6, bcTypes);
    opHier.push_back(pOp);
    if(n != 0) {
      EmIntergrid<Dim> *pT = new EmIntergrid<Dim>(*pGrid, *(gridHier[n-1]));
      trHier.push_back(pT);
    }
  }

  auto PHI_ = [&](rVec a) {
    a = a - translate;
    rVec x { cos(theta) * a[0] + sin(theta) * a[1],
             -sin(theta) * a[0] + cos(theta) * a[1] };
    return PHI(x);
  };
  auto D2PHI_ = [&](rVec a) {
    a = a - translate;
    rVec x { cos(theta) * a[0] + sin(theta) * a[1],
             -sin(theta) * a[0] + cos(theta) * a[1] };
    return D2PHI(x);
  };
  std::vector<std::array<Real, 3>> errnorm(numCompHier);
  for(int n = numCompHier-1; n >= 0; --n) {
    std::vector<const MGLevelOp<2> *> _opHier(opHier.cbegin() + n, opHier.cend());
    std::vector<const Intergrid<2> *> _trHier(trHier.cbegin() + n, trHier.cend());
    std::cout << "\n\nTesting h = " << oneOver(gridSize[n]) << std::endl;
    errnorm[n] = testMixedDer(_opHier, _trHier, PHI_, D2PHI_);
  }

  for(int n = numGrid-1; n >= 0; --n) {
    if(n != 0)
      delete trHier[n-1];
    delete opHier[n];
    delete gridHier[n];
  }
  printConvergenceTable(&gridSize[0], errnorm);
}

int main()
{
  reset_dbglevel(1);
  A[0] = B[0] * ipow<2>(cos(theta)) - B[1] * sin(theta) * cos(theta) + B[2] * ipow<2>(sin(theta));
  A[1] = (B[0] - B[2]) * sin(2 * theta) + B[1] * cos(2 * theta);
  A[2] = B[0] * ipow<2>(sin(theta)) + B[1] * sin(theta) * cos(theta) + B[2] * ipow<2>(cos(theta));
  std::cout << "A = " << A << "\n";
  std::cout << "============================================================================\n";
  doTest({512, 256, 128, 64, 32, 16, 8},
         4,
         [](Vec<Real, 2> x) {
           return sin(2 * M_PI * x[0]) * cos(2 * M_PI * x[1]);
         },
         [&](Vec<Real, 2> x) {
           Real xx = (-4*M_PI*M_PI) * sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]);
           Real xy = (-4*M_PI*M_PI) * cos(2*M_PI*x[0]) * sin(2*M_PI*x[1]);
           Real yy = (-4*M_PI*M_PI) * sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]);
           return B[0]*xx + B[1]*xy + B[2]*yy;
         });
  std::cout << std::endl << std::endl;
  return 0;
}

#else

int main()
{
  std::cout << "============================================================================\n";
  std::cout << "testMixedDerSol is skipped for 3D. " << std::endl;
  return 0;
}

#endif // DIM == 2
